#!/bin/bash  
  
echo "-----Start to set local repo-------"

# 1. 解压gpu_tools_repo.tar.gz到/root目录  
tar -xvf gpu_tools_repo.tar.gz -C /root  
  
# 2. 启动本地更新源服务（使用sealos）  
# 注意：这里假设您已经安装了sealos，并且它可以访问labring/gpu-tools-repo:v2.4镜像  
# 同时也假设32333端口是本地更新源服务的端口  
# 给k8s中的sealos部署节点 打上lable "app.service=gpurepo-node"，后面更新源服务器会根据节点lable调度到sealos 部署节点上
for i in `hostname`
do
  kubectl label nodes $i app.service=gpurepo-node
done


sealos run labring/gpu-tools-repo:v2.4 -e repopath="/root/gpu_tools_repo"  


# 3. 配置当前节点的更新源配置  
hostname=`hostname`
hostip=`cat /etc/hosts | grep $hostname | awk '{ print $1 }'`

if [ x"$hostip" == x"" ]; then
  echo "Can not get hostip in /etc/hosts. Please add 'hostip hostname' peer in /etc/hosts file"
  exit 2
fi


cat <<EOF > /root/drivers-tools.repo
[drivers]
name=Drivers
#baseurl=http://${hostip}:32333/nestos/2203sp3/\$basearch/
baseurl=http://${hostip}:32333/
metadata_expire=1h
enabled=1
gpgcheck=0

EOF
  
# 4. 假设您有一个集群，并且sealos可以访问所有节点  
# 使用sealos将更新源配置文件复制到集群内所有节点上  
  
echo "Copying to nodes"  
sealos scp "/root/drivers-tools.repo" "/etc/yum.repos.d/drivers.repo"  


echo "---------start to install driver--------"

# 5. 安装GPU卡驱动
# 函数：重启服务器
reboot_allhost(){
    CLUSTERFILE="/root/.sealos/default/Clusterfile"
    hostname=`hostname`
    IP1=`cat /etc/hosts | grep $hostname | awk '{ print $1 }'`

    # 使用awk提取IP地址
    IPs=$(awk '/^[[:blank:]]*- ips:$/ { in_ips=1; next }
     in_ips && /^[[:blank:]]*- [0-9]+\.[0-9]+\.[0-9]+\.[0-9]+:[0-9]+$/ {
         split($0, a, ":");
         print a[1];
     }
     /^[[:blank:]]*[^-].*:$/ { in_ips=0 }' "$CLUSTERFILE")

    # 检查提取的IPs中是否包含IP1

    for ip in $IPs; do
        if [ "$ip" != "$IP1" ]; then
             sealos exec --ips $ip "reboot" 2> /dev/null
        fi
    done
    sleep 10
    reboot

}

# 函数：安装海飞科驱动
install_hexaflake_driver() {  
    sealos exec "rpm-ostree install kmod-hexaflake_driver"  
    echo "海飞科驱动安装完成，是否需要重启服务器以应用更改？"  
    read -p "请输入'y'或'n'： " choice  
    case "$choice" in  
        [Yy]* ) reboot_allhost ;;  
        [Nn]* ) echo "请记得稍后手动重启服务器以使驱动生效。" ;;  
        * ) echo "无效输入";;  
    esac  
}  
  
# 函数：安装英伟达驱动和配置  
install_nvidia_driver() { 
    OS=$(detect_node_os)
    if [ x"$OS" == x"" ]; then
      error "Unable to detect OS on node."
      exit 2
    fi

    case $OS in
    *"nestos"* )
      sealos exec "rpm-ostree install kmod-nvidia_driver libnvidia-container-devel libnvidia-container-tools libnvidia-container1 nvidia-container-runtime nvidia-container-toolkit nvidia-container-toolkit-base"  
      ;;
    *"kylin"* | *"openeuler"* )
      sealos exec "yum install -y libnvidia-container-devel libnvidia-container-tools libnvidia-container1 nvidia-container-runtime nvidia-container-toolkit nvidia-container-toolkit-base"  
      ;;
    *)
      warn "Stop the driver installer because '$OS' is not supported in the current version."
      exit 1
      ;;
     esac


  
    # 生成支持nvidia GPU卡的containerd配置文件  
    cat > /root/nvidia-containerd-config.toml <<- EOF2
version = 2
root = "/var/lib/containerd"
state = "/run/containerd"
oom_score = 0

[grpc]
  address = "/run/containerd/containerd.sock"
  uid = 0
  gid = 0
  max_recv_message_size = 16777216
  max_send_message_size = 16777216

[debug]
  address = "/run/containerd/containerd-debug.sock"
  uid = 0
  gid = 0
  level = "warn"

[timeouts]
  "io.containerd.timeout.shim.cleanup" = "5s"
  "io.containerd.timeout.shim.load" = "5s"
  "io.containerd.timeout.shim.shutdown" = "3s"
  "io.containerd.timeout.task.state" = "2s"

[plugins]
  [plugins."io.containerd.grpc.v1.cri"]
    sandbox_image = "sealos.hub:5000/pause:3.8"
    max_container_log_line_size = -1
    max_concurrent_downloads = 20
    disable_apparmor = true
    [plugins."io.containerd.grpc.v1.cri".containerd]
      snapshotter = "overlayfs"
      default_runtime_name = "nvidia"
      [plugins."io.containerd.grpc.v1.cri".containerd.runtimes]
        [plugins."io.containerd.grpc.v1.cri".containerd.runtimes.nvidia]
          runtime_type = "io.containerd.runc.v2"
          runtime_engine = ""
          runtime_root = ""
          [plugins."io.containerd.grpc.v1.cri".containerd.runtimes.nvidia.options]
            BinaryName = "/usr/bin/nvidia-container-runtime"
            SystemdCgroup = true
    [plugins."io.containerd.grpc.v1.cri".registry]
      config_path = "/etc/containerd/certs.d"
      [plugins."io.containerd.grpc.v1.cri".registry.configs]
          [plugins."io.containerd.grpc.v1.cri".registry.configs."sealos.hub:5000".auth]
            username = "admin"
            password = "passw0rd"
  
EOF2
  
    # 备份原始配置文件（如果存在）  
    sealos exec "cp /etc/containerd/config.toml /etc/containerd/config.toml.bak" 
    # 更新配置文件
    sealos scp "/root/nvidia-containerd-config.toml" "/etc/containerd/config.toml"

    # 重启conterind服务
    sealos exec "systemctl restart containerd"
    #安装Nvidia GPU设备插件
    sealos run labring/nvidia-k8s-device-plugin:v0.15.0 -e sharenum=2

    echo "英伟达驱动安装完成，是否需要重启服务器以应用更改？"
    read -p "请输入'y'或'n'： " choice
    case "$choice" in
        [Yy]* ) reboot_allhost ;;
        [Nn]* ) echo "请记得稍后手动重启服务器以使驱动生效。" ;;
        * ) echo "无效输入";;
    esac
}  
 
detect_node_os()
{

  OS=$(grep -E "^ID=" /etc/os-release | cut -d= -f2)
  if [[ -z "${OS}" ]]; then
    OS=$(grep -E "^ID_LIKE=" /etc/os-release | cut -d= -f2)
  fi
  echo "$OS"
}



# 主菜单  
while true; do  
    clear  
    echo "请选择要安装的驱动和工具包："  
    echo "1. 安装海飞科驱动"  
    echo "2. 安装英伟达驱动和配置"  
    echo "q. 退出"  
    read -p "请输入选项（1/2/q）： " option  
    case "$option" in  
        1 ) install_hexaflake_driver; break ;;  
        2 ) install_nvidia_driver; break ;;  
        q|Q ) exit 0 ;;  
        * ) echo "无效选项，请重新输入。";;  
    esac  
done  

# 脚本结束提示  
echo "Update source configuration completed."
